import base64
import datetime
import io
import json
import uuid
from io import BytesIO

from PIL.Image import Image
from requests_toolbelt import MultipartEncoder

from app.project.doc_to_recommendation.llm.model.HttpChat_base_model import HttpChatBaseModel
from app.project.doc_to_recommendation.llm.register.llm_register import LLM_REGISTER
from app.project.doc_to_recommendation.utils.log_utils import get_logger
import requests

@LLM_REGISTER.register_model("http_chat_img_model")
class HttpChatImgModel(HttpChatBaseModel):

    def __init__(self, config: dict):
        self.upload_add_prps = {}
        super().__init__(config)
        self.url = config.get('upload_url', 'http://20.240.22.74:8000/api/common/file/upload')
        self.app_id = config.get('app_id', '67ea5dd3b651523592959985')
        self.bucket_name = config.get('bucket_name', 'chat')
        self.apikey = config.get('apikey', 'gbNK6PbLwygzRlFNc6YXMRhzhukkYQ9pJQqPDNEogGrbpkJLMKwy341mABONG4')

    def fill_header_prps(self):
        super().fill_header_prps()
        upload_header_add_prps = self.config.get('upload_header_prps')
        if upload_header_add_prps is not None:
            self.upload_add_prps.update(upload_header_add_prps)


    def bind_tools_df(self):
        pass

    def agent_calls(self, text, image, prompt=None):
        logger = get_logger(__name__)
        # 上传文件
        # url = self.upload_file(image)

        base64_data = ''
        with open(image['fpath_name'], 'rb') as image_file:
            image_data = image_file.read()
            base64_data = base64.b64encode(image_data).decode('utf-8')

        # 调llm解析图片中文含义
        img_result = super().agent_calls(text, base64_data, prompt)
        if self.bind_t:
            if len(img_result.tool_calls) >= 1:
                imgDscp = img_result.tool_calls[0]['args']
                if imgDscp['type'] == 'img':
                    return {'name': text, 'type': imgDscp['type'], 'description': imgDscp['description']}
            else:
                print("figure:", text, " type:", img_result)
                return {'name': text, 'type': 'img', 'description': img_result.content}
        else:
            content = img_result.get('txBody', {}).get('txEntity', {}).get('choices', [{}])[0].get('message', {}).get('content', '')
            res = content.replace('json', '', 1).replace('\n', '').replace('`', '')
            logger.info("image res:{}".format(res))
            if len(res) == 0:
                return {'type': 'ppt'}
            try:
                res = json.loads(res)
            except:
                logger.warning("json loads error!!!default return as ppt")
                return {'type': 'ppt'}
            return {'name': text, 'type': res['type'], 'description': res['description']}

    def upload_file(self, image) -> str:
        # 生成唯一ID和跟踪号
        chat_id = str(uuid.uuid4())
        current_time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")

        global_track_no = f"{current_time}997140214300752110"

        base64_data = ''
        with open(image['fpath_name'], 'rb') as image_file:
            image_data = image_file.read()
            base64_data = base64.b64encode(image_data).decode('utf-8')

        # 可完全自定义的请求体
        payload = {
            "txHeader": {
                "servNo": "rag_v1_C011003",
                "startSysOrCmptNo": "99714020000",
                "sendSysOrCmptNo": "99714020000",
                "startChnlFgCd": "01",
                "busiSendInstNo": "11005293",
                "dataCenterCode": "Y",
                "txStartTime": "20201103202020001",
                "txSendTime": "20201103202020001",
                "msgrptTotalLen": "999",
                "msgrptFmtVerNo": "00001",
                "msgAgrType": "T",
                "pubMsgHeadLen": "999",
                "embedMsgrptLen": "999",
                "targetSysOrCmptNo": "99900180002",
                "servTpCd": "1",
                "servVerNo": "00001",
                "globalBusiTrackNo": global_track_no,
                "subtxNo": "17143007518028579764456528793235",
                "reqSysSriNo": "2021077317321202021200123",
                "resvedInputInfo": "info"
            },
            "txBody": {
                "txEntity": {
                    "fileBase64": base64_data,
                    "fileName": image['fpath_name'],
                    "bucketName": self.bucket_name,
                    "metadata": json.dumps({"chatId": chat_id}),
                    "data": json.dumps({"appId": self.app_id}),
                    "bmdlJoinupSceneNo": "997140200000355"
                }
            }
        }

        print('============upload_file post============', self.url, payload)
        headers = {"Content-Type": "application/json"}
        headers.update(self.upload_add_prps) if self.upload_add_prps is not None else None
        response = requests.post(self.url, json=payload, headers=headers)
        print('==========upload_file response=============', response.json())
        data = response.json()
        return data.get('txBody', {}).get('txEntity', {}).get('data', {}).get('previewUrl')
